import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import os


def get_transforms(train=True, img_size=224, mean=None, std=None):
    """
    获取数据预处理变换
    
    Args:
        train (bool): 是否为训练模式
        img_size (int): 图像大小
        mean (list): 标准化均值
        std (list): 标准化标准差
    
    Returns:
        transforms.Compose: 数据变换组合
    """
    if mean is None:
        mean = [0.485, 0.456, 0.406]
    if std is None:
        std = [0.229, 0.224, 0.225]
    
    if train:
        transform = transforms.Compose([
            # transforms.Resize((img_size, img_size)),
            transforms.RandomResizedCrop(img_size),
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.RandomRotation(10),
            transforms.ToTensor(),
            transforms.Normalize(mean=mean, std=std)
        ])
    else:
        transform = transforms.Compose([
            transforms.Resize((img_size, img_size)),
            transforms.ToTensor(),
            transforms.Normalize(mean=mean, std=std)
        ])
    
    return transform


def create_data_loaders(data_dir, batch_size=256, num_workers=0, pin_memory=True, 
                       img_size=224, mean=None, std=None):
    """
    创建训练、验证和测试数据加载器
    
    Args:
        data_dir (str): 数据集根目录
        batch_size (int): 批次大小
        num_workers (int): 数据加载器工作进程数
        pin_memory (bool): 是否使用pin_memory
        img_size (int): 图像大小
        mean (list): 标准化均值
        std (list): 标准化标准差
    
    Returns:
        tuple: (train_loader, val_loader, test_loader, train_dataset)
    """
    # 数据集路径
    train_dir = os.path.join(data_dir, 'train')
    val_dir = os.path.join(data_dir, 'val')
    test_dir = os.path.join(data_dir, 'test')
    
    # 获取数据变换
    train_transform = get_transforms(train=True, img_size=img_size, mean=mean, std=std)
    val_test_transform = get_transforms(train=False, img_size=img_size, mean=mean, std=std)
    
    # 加载数据集
    train_dataset = datasets.ImageFolder(train_dir, transform=train_transform)
    val_dataset = datasets.ImageFolder(val_dir, transform=val_test_transform)
    test_dataset = datasets.ImageFolder(test_dir, transform=val_test_transform)
    
    # 创建数据加载器
    train_loader = DataLoader(
        train_dataset,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=pin_memory
    )
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=pin_memory
    )
    
    test_loader = DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=pin_memory
    )
    
    return train_loader, val_loader, test_loader, train_dataset 